from typing import List, Union

from pymodels.utils.BookReader import BasicBookReader
from pymodels.models.dpr import tinybertVectorizer
from pymodels.utils.Indexer import FaissIndexerHelper
from pluginUtils import BasicReaderPlugin

from transformers import BertConfig


class ReaderPlugin(BasicReaderPlugin, BasicBookReader):
    def __init__(self, modelname: str = "zerohell/tinydpr-acc_0.315-bs_307", device: str = "cpu",
                 size: int = 512, encoding: str = 'utf8'):
        """
        设置Txt阅读器。
        :param modelname: DPR模型路经。
        :param device: 推理时所使用的处理器。'cpu'或'cuda:n'。推荐采用：cuda:0
        :param size: 读取文本时的每页的大小。越小，推理器的结果可能越精确。推荐值：512.
        :param encoding: 读取文件时采用的编码。
        :return:
        """
        self.modelname = modelname
        self.device = device
        dim = BertConfig.from_pretrained(modelname).hidden_size
        indexer = FaissIndexerHelper(dim=dim)
        self.page_size = 0
        self.size = size
        self.encoding = encoding

        BasicBookReader.__init__(self, indexer=indexer, booktype='.txt', cache=True, arduino=True)
        BasicReaderPlugin.__init__(self)

    def hibernate(self):
        BasicBookReader.del_vectorizer(self)

    def wake(self):
        dpr = tinybertVectorizer(self.modelname)
        dpr.model.to(self.device)
        BasicBookReader.set_vectorizer(self, dpr)

    def file2pages(self, filename: str, mode='dense') -> List[str]:
        """
        :param filename: 要处理的文件名称。
        :param mode: 设置文件的转换数格式。`dense`将文本视为连续的字符流。
        :return: 页形成的数组。会包括一个空页，用于判断是否为不相关的问题。
        """
        size = self.size
        encoding = self.encoding
        blocks = ['']
        overlap = size // 2
        self.page_size = size
        with open(filename, encoding=encoding, mode='r') as file:
            if mode == 'dense':
                remainedline = ''
                for line in file.readlines():
                    remainedline += line
                    while len(remainedline) >= size:
                        blocks.append(remainedline[:size])
                        remainedline = remainedline[overlap:]
                if len(remainedline) > 0:
                    blocks.append(remainedline)
            else:
                raise Exception("未实现的模式。")
        return blocks

    def readbook(self, bookpath: str):
        return BasicBookReader.readbook(self, filename=bookpath)

    def get_readfile_name(self):
        return BasicBookReader.get_readfile_name(self)

    def search(self, query: Union[str, List[str]], k: int = 1):  # noqa
        return BasicBookReader.search(self, query, k=k, return_page_content=True, return_cos_similarity=True)

    def get_page_size(self) -> int:
        return self.page_size

    def get_introduction(self):
        return """
        作者：xzh
        日期：2023年3月26日
        邮箱：kodderopert@163.com
        """
